//
// Created by lanhu on 2024/8/17.
//

#include <format>

#include "ncnn_instance.h"
#include "print_log.h"


static inline float intersection_area(const Object& a, const Object& b)
{
    cv::Rect_<float> inter = a.rect & b.rect;
    return inter.area();
}

static void nms_sorted_bboxes(const std::vector<Object>& faceObjects, std::vector<int>& picked, float nms_threshold)
{
    picked.clear();

    const int n = faceObjects.size();

    std::vector<float> areas(n);
    for (int i = 0; i < n; i++)
    {
        areas[i] = faceObjects[i].rect.area();
    }

    for (int i = 0; i < n; i++)
    {
        const Object& a = faceObjects[i];

        bool keep = true;
        for (const int pick : picked)
        {
            const Object& b = faceObjects[pick];

            if (a.label != b.label)
            {
                continue;
            }
            float inter_area = intersection_area(a, b);
            float union_area = areas[i] + areas[pick] - inter_area;
            if (inter_area / union_area > nms_threshold)
            {
                keep = false;
            }

        }

        if (keep)
        {
            picked.push_back(i);
        }
    }
}

static inline float sigmoid(float x)
{
    return static_cast<float>(1.f / (1.f + exp(-x)));
}

static inline float clampf(float d, float min, float max)
{
    const float t = d < min ? min : d;
    return t > max ? max : t;
}

static void parse_detections(float* inputs, float confidence_threshold, int num_channels, int num_anchors,
                             int num_labels, int infer_img_width, int infer_img_height, std::vector<Object>& objects)
{
    std::vector<Object> detections;
    cv::Mat output = cv::Mat((int)num_channels, (int)num_anchors, CV_32F, inputs).t();

    for (int i = 0; i < num_anchors; i++)
    {
        auto row_ptr = output.row(i).ptr<float>();
        auto bboxes_ptr = row_ptr;
        auto scores_ptr = row_ptr + 4;
        auto max_s_ptr = std::max_element(scores_ptr, scores_ptr + num_labels);
        float score = *max_s_ptr;
        if (score > confidence_threshold)
        {
            float x = *bboxes_ptr++;
            float y = *bboxes_ptr++;
            float w = *bboxes_ptr++;
            float h = *bboxes_ptr;

            float x0 = clampf((x - 0.5f * w), 0.f, (float)infer_img_width);
            float y0 = clampf((y - 0.5f * h), 0.f, (float)infer_img_height);
            float x1 = clampf((x + 0.5f * w), 0.f, (float)infer_img_width);
            float y1 = clampf((y + 0.5f * h), 0.f, (float)infer_img_height);

            cv::Rect_<float> bbox;
            bbox.x = x0;
            bbox.y = y0;
            bbox.width = x1 - x0;
            bbox.height = y1 - y0;
            Object object;
            object.label = max_s_ptr - scores_ptr;
            object.prob = score;
            object.rect = bbox;
            detections.emplace_back(object);
        }
    }
    objects = detections;
}



NCNNInstance::NCNNInstance(): m_net(nullptr) {

}
int NCNNInstance::Init(const char* modelDir) {
    if (m_net == nullptr) {
        m_net = new ncnn::Net();
    }

    if (modelDir == nullptr) {
        LOGE("NCNNInstance::Init fail! param is null!");
        return -1;
    }
    ncnn::create_gpu_instance();

    m_net->opt.use_fp16_packed = false;
    m_net->opt.use_fp16_storage = false;
    m_net->opt.use_fp16_arithmetic = false;
    m_net->opt.num_threads = 4;
    m_net->opt.use_vulkan_compute = false;
    m_net->opt.lightmode = true;
    m_net->opt.blob_allocator = &(m_blob_pool_allocator);
    m_net->opt.workspace_allocator = &(m_workspace_pool_allocator);


    char buf[100];
    std::sprintf(buf, "%s/model/ncnn.param", modelDir);
    if(0 != m_net->load_param(buf)) {
        LOGE("NCNNInstance::Init fail! can't load model param file!");
        return -1;
    }
    std::sprintf(buf, "%s/model/ncnn.bin", modelDir);
    if(0 != m_net->load_model(buf)) {
        LOGE("NCNNInstance::Init fail! can't load model bin file!");
        return -1;
    }
    return 0;
}

int NCNNInstance::Detect(const cv::Mat& image) {
    if (m_net == nullptr) {
        return -1;
    }

    const int imageWidth = image.cols;
    const int imageHeight = image.rows;

    m_objects.clear();
    m_proposals.clear();

    const float scale = static_cast<float>(DETECT_SUB_SIZE) / imageWidth;
    const int w = DETECT_SUB_SIZE;
    const int h = imageHeight * scale;


    ncnn::Mat in = ncnn::Mat::from_pixels_resize(image.data, ncnn::Mat::PIXEL_BGRA2BGR, imageWidth, imageHeight, DETECT_SUB_SIZE, h);

    const int wPad = (DETECT_SUB_SIZE + MAX_STRIDE - 1) / MAX_STRIDE * MAX_STRIDE - w;
    const int hPad = (DETECT_SUB_SIZE + MAX_STRIDE - 1) / MAX_STRIDE * MAX_STRIDE - h;
    ncnn::Mat inPad;
    copy_make_border(in, inPad, hPad / 2, hPad - hPad / 2, wPad / 2, wPad - wPad / 2, ncnn::BORDER_CONSTANT, 114.f);

    constexpr float normVals[3] = { 1 / 255.f, 1 / 255.f, 1 / 255.f };
    inPad.substract_mean_normalize(0, normVals);


    ncnn::Extractor ex = m_net->create_extractor();
    const int ret = ex.input("in0", inPad);
    if (ret != 0) {
        return -1;
    }

    ncnn::Mat out;
    ex.extract("out0", out);

    parse_detections(static_cast<float*>(out.data), PROB_THRESHOLD, out.h, out.w, LABEL_NUM, inPad.w, inPad.h,m_proposals);

    std::vector<int> picked;
    nms_sorted_bboxes(m_proposals, picked, NMS_THRESHOLD);

    m_objects.resize(picked.size());
    for (int i = 0; i < picked.size(); i++)
    {
        m_objects[i] = m_proposals[picked[i]];

        float x0 = (m_objects[i].rect.x - (wPad / 2)) / scale;
        float y0 = (m_objects[i].rect.y - (hPad / 2)) / scale;
        float x1 = (m_objects[i].rect.x + m_objects[i].rect.width - (wPad / 2)) / scale;
        float y1 = (m_objects[i].rect.y + m_objects[i].rect.height - (hPad / 2)) / scale;

        x0 = std::max(std::min(x0, (float)(imageWidth - 1)), 0.f);
        y0 = std::max(std::min(y0, (float)(imageHeight - 1)), 0.f);
        x1 = std::max(std::min(x1, (float)(imageWidth - 1)), 0.f);
        y1 = std::max(std::min(y1, (float)(imageHeight - 1)), 0.f);

        m_objects[i].rect.x = x0;
        m_objects[i].rect.y = y0;
        m_objects[i].rect.width = x1 - x0;
        m_objects[i].rect.height = y1 - y0;
        // LOGD("label:%d x:%d y%d w:%d h:%d", m_objects[i].label, m_objects[i].rect.x, m_objects[i].rect.y, m_objects[i].rect.width, m_objects[i].rect.height);
    }
    return 0;
}

void NCNNInstance::Destroy() {
    if (m_net != nullptr) {
        ncnn::destroy_gpu_instance();
        m_net->clear();
        delete m_net;
        m_net = nullptr;
    }
}

NCNNInstance::~NCNNInstance() {
    Destroy();
}
